"""
AI 服务基类
定义不同 AI 功能的抽象接口
"""

from abc import ABC, abstractmethod
from typing import Any, Dict, List, Optional

from pydantic import BaseModel


class AIResponse(BaseModel):
    """AI 响应基类"""

    success: bool
    provider: str
    model: str
    feature_type: str
    usage: Optional[Dict[str, Any]] = None
    cost: Optional[float] = None
    metadata: Optional[Dict[str, Any]] = None


class TextGenerationResponse(AIResponse):
    """文本生成响应"""

    feature_type: str = "text-generation"
    content: str
    finish_reason: Optional[str] = None


class TextToImageResponse(AIResponse):
    """文生图响应"""

    feature_type: str = "text-to-image"
    images: List[str]  # 图片 URL 列表
    revised_prompt: Optional[str] = None  # 修正后的提示词


class ImageToTextResponse(AIResponse):
    """图生文响应"""

    feature_type: str = "image-to-text"
    content: str
    confidence: Optional[float] = None


class BaseAIService(ABC):
    """AI 服务基类"""

    def __init__(self, provider: str, model: str, config: Dict[str, Any]):
        self.provider = provider
        self.model = model
        self.config = config

    @abstractmethod
    async def execute(self, *args, **kwargs) -> AIResponse:
        """执行 AI 请求"""
        pass

    def calculate_cost(self, usage: Dict[str, Any]) -> float:
        """计算成本"""
        return 0.0


class TextGenerationService(BaseAIService):
    """文本生成服务基类"""

    @abstractmethod
    async def generate(
        self,
        prompt: str,
        system_prompt: Optional[str] = None,
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
        stream: bool = False,
        **kwargs,
    ) -> TextGenerationResponse:
        """生成文本"""
        pass

    async def execute(
        self,
        prompt: str,
        system_prompt: Optional[str] = None,
        **kwargs,
    ) -> TextGenerationResponse:
        """执行文本生成"""
        return await self.generate(prompt, system_prompt, **kwargs)


class TextToImageService(BaseAIService):
    """文生图服务基类"""

    @abstractmethod
    async def generate_image(
        self,
        prompt: str,
        width: Optional[int] = None,
        height: Optional[int] = None,
        num_images: int = 1,
        style: Optional[str] = None,
        **kwargs,
    ) -> TextToImageResponse:
        """生成图片"""
        pass

    async def execute(
        self,
        prompt: str,
        **kwargs,
    ) -> TextToImageResponse:
        """执行文生图"""
        return await self.generate_image(prompt, **kwargs)


class ImageToTextService(BaseAIService):
    """图生文服务基类"""

    @abstractmethod
    async def analyze_image(
        self,
        image_url: str,
        prompt: Optional[str] = None,
        **kwargs,
    ) -> ImageToTextResponse:
        """分析图片"""
        pass

    async def execute(
        self,
        image_url: str,
        **kwargs,
    ) -> ImageToTextResponse:
        """执行图生文"""
        return await self.analyze_image(image_url, **kwargs)

